package class02

import (
	"math/rand"
	"time"
)

/*
	长度为N的数组arr，一定可以组成N^2个数值对
	例如：arr = [3,1,2]
	数值对有（3，3）（3，1）（3，2）（1，3）（1，1）（1，2）（2，3）（2，1）（2，2）
	也就是任意两个数都有数值对，而且自己和自己也算数值对
	数值对怎么排序？规定，第一维数据从小到大排序，第一维数据一样的，第二维数据也从小到大，所以
	上面的数值对排序的结果为：
	（1，1）（1，2）（1，3）（2，1）（2，2）（2，3）（3，1）（3，2）（3，3）
	给定一个数组arr，和整数K，返回第k小的数值对。
*/

func kthMinPair(arr []int, k int) []int {
	size := len(arr)
	ans := make([]int, 2)
	firstIndex := (k-1)/size
	ans[0] = GetMinKth(arr, firstIndex)
	lessNum := 0
	firstNumSize := 0
	for i:=0; i<len(arr); i++ {
		if arr[i] < ans[0] {
			lessNum++
		}
		if arr[i] == ans[0] {
			firstNumSize++
		}
	}
	rest := k-lessNum*size
	ans[1] = GetMinKth(arr, (rest-1)/firstNumSize)
	return ans
}

// 再arr中找到第k小的数
// k可以冲0开始
// 返回值第二个参数表示相同的数有几个
func GetMinKth(arr []int, k int) int {
	l := 0
	r := len(arr)-1
	var pivot int // 支点，随机寻找数组上的点做partition
	var rl, rr int
	rand.Seed(int64(time.Now().Nanosecond()))
	for l < r {
		pivot = arr[rand.Intn(r-l+1) + l]
		rl, rr = partition(arr, l, r, pivot)
		if k < rl {
			r = rl-1
		} else if k > rr {
			l = rr+1
		} else {
			return pivot
		}
	}
	return arr[l]
}

func partition(arr []int, l, r, pivot int) (int, int) {
	less := l-1
	more := r+1
	cur := l
	for cur < more {
		if arr[cur] < pivot {
			less++
			swap(arr, less, cur)
			cur++
		} else if arr[cur] > pivot {
			more--
			swap(arr, more, cur)
		} else {
			cur++
		}
	}
	return less+1, more-1
}

func swap(arr []int, index1, index2 int) {
	arr[index1], arr[index2] = arr[index2], arr[index1]
}